import os
os.getcwd()
from model.vgg_fp2a_3cls import vgg16_bn as fp2a
import os
import torch
from PIL import Image
from torchvision import transforms


def test(img_path, model, device):
    mean, std = [0.4914, 0.4822, 0.4465], [0.2023, 0.1994, 0.2010]
    _transforms = transforms.Compose([transforms.Resize(256),
                                        transforms.CenterCrop(224),
                                        transforms.ToTensor(),
                                        transforms.Normalize(mean, std)])
    assert os.path.exists(img_path), "file: '{}' dose not exist.".format(img_path)
    img = Image.open(img_path)
    img = _transforms(img)
    img = torch.unsqueeze(img, dim=0).to(device)
    model.eval()
    # predict class
    outputs = model(img, PATH=img_path)
    predict = torch.softmax(outputs, dim=0)
    predict_cla = torch.argmax(predict).cpu().numpy()
    print('predicted result: ', predict_cla)


def load_model(model, state_dict, device):
    model.load_state_dict(state_dict)
    model.to(device)
    model.eval()
    return model


if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    net1 = fp2a()
    params = torch.load('./checkpoint/best_acc_160.pth')
    model = load_model(net1, params, device)
    jpg_path = "./dataset/test.jpg"

    test(jpg_path, model, device)
